from math import ceil

import numpy as np
import pandas as pd
import seaborn as sns
import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.patches as mpatches
import scipy as stats
from adjustText import adjust_text
from scipy.cluster.hierarchy import linkage 
from scipy.cluster.hierarchy import dendrogram
from matplotlib.legend_handler import HandlerTuple


def dep_bar_plot(
    reg_count_result: OrderedDict[str, dict[str, int]], 
    output_filename: str
) -> None:
    
    n_group = len(reg_count_result)
    figure_size = (3 * n_group, 5)
    top_margin = 2
    bar_width, half_width = 0.2, 0.1
    max_n_yticks = 10

    fig, ax = plt.subplots(figsize=figure_size)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.set_xlim([0.5, n_group + 0.5])
    ax.set_xticks(
        list(range(1, n_group + 1)), 
        list(reg_count_result.keys()), 
        rotation=30, ha='right'
    )
    ax.set_ylabel("Number of Proteins")
    ax.set_title(
        "Differentially expressed proteins in different groups", pad=20
    )

    bars = list()
    highest_count = 0
    first_pair = True
    for x, ctr in enumerate(reg_count_result.values(), start=1):
        ups, downs = ctr["up"], ctr["down"]
        highest_count = (
            v if (v := max(ups, downs)) > highest_count else highest_count
        )
        if first_pair:
            up_label, down_label = "up", "down"
            first_pair = False
        else:
            up_label, down_label = None, None
            
        b1, = ax.bar(
            x - half_width, ups, label=up_label, 
            width=bar_width, color="darkred"
        )
        b2, = ax.bar(
            x + half_width, downs, label=down_label, 
            width=bar_width, color="darkblue"
        )
        bars.append(b1)
        bars.append(b2)

    n_yticks = min(max_n_yticks, highest_count)
    tick_step = ceil(highest_count / n_yticks)
    y_top_lim = highest_count + top_margin * tick_step
    ax.set_ylim([0, y_top_lim])
    ax.set_yticks(
        list(range(0, highest_count + top_margin * tick_step + 1, tick_step))
    )
    ax.legend()

    for bar in bars:
        height = bar.get_height()
        ax.text(
            bar.get_x() + bar.get_width() / 2, 
            height, f'{height:.0f}', 
            ha='center', va='bottom'
        )

    plt.savefig(output_filename, dpi=300)


def volcano_plot(
    de_results: pd.DataFrame, 
    output_filename: str, 
    fdr_threshold: float=0.05
) -> None:

    plt.figure(figsize=(6, 6))
    volcano = sns.scatterplot(
        data=de_results, 
        x="log2FC", y=-np.log10(de_results["fdr"]),
        hue="regulation",
        palette={
            "Sig_Up": "red", 
            "Sig_Down": "blue", 
            "NoDiff": "grey"
        },
        alpha=0.7, s=30
    )

    plt.axhline(
        -np.log10(fdr_threshold), linestyle='--', color='grey', alpha=0.7
    )
    plt.axvline(0, linestyle='--', color='grey', alpha=0.7)

    plt.xlabel("log2(FC)")
    plt.ylabel("-log10(FDR)")
    plt.title("Advanced Volcano Plot", pad=20)
    plt.grid(True, linestyle='--', alpha=0.2)

    plt.legend(title='Significant', loc='best')

    sig_df = de_results[de_results["significant"]]
    texts = []
    for prot, row in sig_df.iterrows():
        texts.append(
            plt.text(
                row['log2FC'], -np.log10(row['fdr']), prot, 
                fontweight="bold", fontsize=8, ha='center', va='bottom'
            )
        )

    adjust_text(texts, arrowprops=dict(arrowstyle='-', color='gray', lw=0.5))
    plt.savefig(output_filename, dpi=300)


def heatmap_plot(
    grouped_npx_dfs: dict[str, pd.DataFrame], 
    de_results: pd.DataFrame, 
    output_filename: str
) -> None:

    sig_assays = de_results[de_results["regulation"] == "Sig_Up"].index
    sig_assays = sig_assays.append(
        de_results[de_results["regulation"] == "Sig_Down"].index
    )
    
    colors = sns.color_palette("husl", len(grouped_npx_dfs))
    grp_color_map = dict()
    grp_n_samples = list()
    grp_labels, npx_dfs, grp_colors = list(), list(), list()
    for idx, (grp, npx_df) in enumerate(grouped_npx_dfs.items()):
        grp_labels.append(grp)
        n_a, n_s = npx_df.shape
        grp_n_samples.append(n_s)
        npx_dfs.append(npx_df)
        c = colors[idx]
        grp_color_map[grp] = c
        grp_colors += [c] * n_s
    full_npx_df = pd.concat(npx_dfs, axis=1).reindex(sig_assays)
    n_assays, n_samples = full_npx_df.shape
    
    fig_h_rate = 8 / 30
    fig_w_rate = 15 / 50
    color_bar_width = 0.15
    group_legend_width = 1.5
    n_groups = len(grouped_npx_dfs)
    fig_w = fig_w_rate * n_samples
    fig_h = max(fig_h_rate * n_assays, 6)
    color_bar_w_ratio = color_bar_width / fig_w
    group_legend_w_ratio = group_legend_width / fig_w
    main_axes_w_ratio = 1 - color_bar_w_ratio - group_legend_w_ratio

    fig, ax = plt.subplots(
        2, 3, 
        width_ratios=(
            main_axes_w_ratio, 
            color_bar_w_ratio, 
            group_legend_w_ratio
        ), 
        height_ratios=(0.05, 0.95), 
        figsize=(fig_w, fig_h)
    )

    group_bar_ax = ax[0, 0]
    group_bar_ax.imshow([grp_colors], aspect="auto")
    group_bar_ax.set_xticks([])    
    group_bar_ax.set_yticks([]) 
    group_bar_ax.spines[:].set_visible(False)

    color_bar_ax = ax[1, 1]
    heatmap_ax = ax[1, 0]
    sns.heatmap(
        full_npx_df, ax=heatmap_ax, 
        cmap="RdYlBu_r", cbar_ax=color_bar_ax, 
        yticklabels=full_npx_df.index, xticklabels=full_npx_df.columns
    )
    heatmap_ax.set_xticklabels(
        heatmap_ax.get_xticklabels(), rotation=45, ha="right", fontsize=9
    )
    heatmap_ax.set(xlabel='', ylabel='')
    for i in grp_n_samples:
        heatmap_ax.axvline(i, color="white", linewidth=2)

    legend_ax = ax[1, 2]
    legend_ax.axis("off")
    patches = [
        mpatches.Patch(color=c, label=g) for g, c in grp_color_map.items()
    ]
    legend_ax.legend(handles=patches, title="group", loc="upper left")

    ax[0, 1].axis("off")
    ax[0, 2].axis("off")
    plt.tight_layout(h_pad=0.2, w_pad=0)
    plt.savefig(output_filename, dpi=300)


def avg_heatmap_plot(
    grouped_npx_dfs: dict[str, pd.DataFrame], 
    de_results: pd.DataFrame,
    output_filename: str
) -> None:  

    sig_assays = de_results[de_results["regulation"] == "Sig_Up"].index
    sig_assays = sig_assays.append(
        de_results[de_results["regulation"] == "Sig_Down"].index
    )
    npx_mean_df = pd.DataFrame(
        {grp:df.reindex(sig_assays).mean(axis=1) 
         for grp, df in grouped_npx_dfs.items()}
    )
    fig_h = max(len(sig_assays) * (2 / 5), 2)
    fig, ax = plt.subplots(
        1, 4, width_ratios=(0.2, 0.5, 0.28, 0.02), figsize=(8, fig_h)
    )
    cluster_tree_ax, heatmap_ax, barplot_ax, colorbar_ax = ax
    
    cluster_tree_ax.axis("off")
    assay_linkage = linkage(npx_mean_df, method="average")    
    assay_dendro = dendrogram(
        assay_linkage, orientation="left", ax=cluster_tree_ax
    )
    
    assay_dendro["leaves"].reverse()
    npx_mean_df = npx_mean_df.iloc[assay_dendro["leaves"], ]
    sns.heatmap(
        npx_mean_df, ax=heatmap_ax, cbar_ax=colorbar_ax,
        yticklabels=npx_mean_df.index, cmap="RdYlBu_r", 
        xticklabels=npx_mean_df.columns, 
    )
    heatmap_ax.set(xlabel='', ylabel='')
    heatmap_ax.set_yticklabels(
        heatmap_ax.get_yticklabels(), fontsize=8
    )

    part_de_results = de_results.reindex(npx_mean_df.index)
    sns.barplot(
        part_de_results, x="log2FC", y=part_de_results.index, ax=barplot_ax
    )
    barplot_ax.set(yticks=[], ylabel="", xlabel="")
    
    plt.tight_layout(w_pad=0.3)
    plt.savefig(output_filename, dpi=300)


def box_plot(
    grouped_npx_dfs: dict[str, pd.DataFrame], 
    de_results: pd.DataFrame, 
    output_filename: str
) -> None:

    sig_assays = de_results[de_results["regulation"] == "Sig_Up"].index
    sig_assays = sig_assays.append(
        de_results[de_results["regulation"] == "Sig_Down"].index
    )

    melt_df_lst = list()
    for grp, df in grouped_npx_dfs.items():
        df = df.reindex(sig_assays).reset_index().melt(
            id_vars="Assay"
        ).rename(columns=dict(value="NPX"))
        df["group"] = grp
        melt_df_lst.append(df)
    df = pd.concat(melt_df_lst)

    bar_width = 0.4

    fig, ax = plt.subplots(figsize=(6, 4))
    sns.boxplot(
        df, x="Assay", y="NPX", hue="group", linewidth=1.2, 
        ax=ax, width=bar_width, fliersize=2.5, order=sig_assays
    )
    ax.tick_params(axis='x', which='major', labelsize=9)
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right")
    ax.set_xlabel("Assay", fontweight="bold", labelpad=6)
    ax.set_ylabel("NPX", fontweight="bold", labelpad=6)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)

    if len(grouped_npx_dfs) == 2:
        g1, g2 = grouped_npx_dfs.keys()
        is_g1, is_g2 = (df["group"] == g1), (df["group"] == g2)
        for idx, assay in enumerate(sig_assays):
            is_that = (df["Assay"] == assay)
            g1_npx_values = df[is_that & is_g1]["NPX"]
            g2_npx_values = df[is_that & is_g2]["NPX"]
            t_stat, p_val = stats.ttest_ind(g1_npx_values, g2_npx_values)
            if p_val < 0.001: symbol = "***"
            elif p_val < 0.01: symbol = "**"
            elif p_val < 0.05: symbol = "*"
            else: symbol = None
            if symbol:
                x1, x2 = idx - bar_width / 2, idx + bar_width / 2
                y = max(g1_npx_values.max(), g2_npx_values.max()) + 0.5
                ax.plot(
                    [x1, x2], [y, y], 
                    lw=1, color='black'
                )
                mid_x = (x1 + x2) / 2
                ax.text(
                    mid_x, y + 0.1, symbol, 
                    ha='center', va='bottom', color='black', 
                    fontsize=12
                )
    
    plt.savefig(output_filename, dpi=300)


def correlation_heatmap(
    grouped_npx_dfs: dict[str, pd.DataFrame], 
    de_results: pd.DataFrame, 
    output_filename: str, 
    star_anno: bool=False
) -> tuple[np.ndarray, np.ndarray]:

    sig_assays = de_results[de_results["regulation"] == "Sig_Up"].index
    sig_assays = sig_assays.append(
        de_results[de_results["regulation"] == "Sig_Down"].index
    )
    sig_assay_all_sample_npx = pd.concat(
        grouped_npx_dfs.values(), axis=1
    ).reindex(sig_assays)

    n_prot, _ = sig_assay_all_sample_npx.shape
    r_mtx = np.ones((n_prot, n_prot))
    p_mtx = np.ones((n_prot, n_prot))
    for i in range(n_prot):
        for j in range(i + 1, n_prot):
            rev = stats.pearsonr(
                sig_assay_all_sample_npx.iloc[i], sig_assay_all_sample_npx.iloc[j]
            )
            r_mtx[i, j], r_mtx[j, i] = rev.statistic, rev.statistic
            p_mtx[i, j], p_mtx[j, i] = rev.pvalue, rev.pvalue
                
    assay_index = sig_assay_all_sample_npx.index
    assay_index.name = None
    data = pd.DataFrame(r_mtx, index=assay_index, columns=assay_index)
    
    if star_anno:        
        annot_matrix = np.full_like(r_mtx, "", dtype=str)
        for i in range(n_prot):
            for j in range(n_prot):
                if i != j and p_mtx[i, j] < 0.001:
                    annot_matrix[i, j] = "***"
                elif i != j and p_mtx[i, j] < 0.01:
                    annot_matrix[i, j] = "**"
                elif i != j and p_mtx[i, j] < 0.05:
                    annot_matrix[i, j] = "*"
        sns.clustermap(
            data, annot=annot_matrix, figsize=(5, 5), fmt='s', 
            cmap="RdYlBu_r", tree_kws={"linewidth":1}
        )
    else:
        sns.clustermap(
            data, annot=True, figsize=(5, 5), 
            cmap="RdYlBu_r", fmt=".2f", tree_kws={"linewidth":1}
        )
    plt.savefig(output_filename, dpi=300)
    # return sig_assay_all_sample_npx
    return (pd.DataFrame(r_mtx, index=assay_index, columns=assay_index), 
            pd.DataFrame(p_mtx, index=assay_index, columns=assay_index))


def correlation_network(
    corr_df: pd.DataFrame, 
    pvalue_df: pd.DataFrame, 
    alpha: float=0.05, 
) -> None:

    base_node_size=50
    node_size_scale=100
    base_edge_width=0.5
    edge_width_scale=4
    node_color = "darkviolet"
    # title="Protein Correlation Network"
    
    input_df_check = (corr_df.index.equals(pvalue_df.index) and 
                      corr_df.columns.equals(pvalue_df.columns) and 
                      corr_df.index.equals(corr_df.columns))
    assert input_df_check
    
    protein_names = corr_df.index.tolist()
    n = len(protein_names)
    r_mtx, p_mtx = corr_df.values, pvalue_df.values
    
    g = nx.Graph()
    for i in range(n):
        for j in range(i + 1, n):
            if p_mtx[i, j] < alpha:
                corr = r_mtx[i, j]
                g.add_edge(
                    protein_names[i], protein_names[j], 
                    corr=corr, 
                    abs_corr=abs(corr), 
                    line_style=("solid" if corr > 0 else "dashed"), 
                    line_color=("goldenrod" if corr > 0 else "gray"), 
                    line_width=base_edge_width + edge_width_scale * abs(corr)
                )
    
    node_degrees = dict(g.degree())
    node_sizes = [
        base_node_size + node_size_scale * node_degrees[nd] for nd in g.nodes()
    ]
    
    fig, (net_ax, deg_leg_ax, rel_leg_ax, abs_leg_ax) = plt.subplots(
        1, 4, figsize=(8, 4), width_ratios=(0.7, 0.1, 0.1, 0.1)
    )
    pos = nx.circular_layout(g)
    nodes = nx.draw_networkx_nodes(
        g, pos, node_size=node_sizes,
        node_color=node_color, ax=net_ax
    )
    net_ax.axis("off")
    deg_leg_ax.axis("off")
    rel_leg_ax.axis("off")
    abs_leg_ax.axis("off")
    
    for u, v in g.edges():
        nx.draw_networkx_edges(
            g, pos, edgelist=[(u, v)],
            width=g[u][v]["line_width"], 
            style=g[u][v]["line_style"],
            edge_color=g[u][v]["line_color"],
            ax=net_ax
        )
        
    nx.draw_networkx_labels(
        g, pos, font_size=10, font_weight="bold", ax=net_ax
    )
    
    rho_values = list(np.arange(0, 1.1, 0.25))
    edge_legend_lines = list()
    for rho in rho_values:
        linewidth = base_edge_width + edge_width_scale * rho
        line = mlines.Line2D(
            [], [], color="black", 
            linewidth=linewidth,
            linestyle='-',
            label=f"{rho:.2f}"
        )
        edge_legend_lines.append(line)
    abs_leg_ax.legend(
        handles=edge_legend_lines, labelspacing=0.5, 
        title="abs(Rho)", frameon=False, framealpha=0.8, 
        title_fontproperties = {"weight":"bold", "size":"medium"}
    )

    pos_legend_line = list()
    pos_line = mlines.Line2D(
        [], [], color="goldenrod", linestyle='-', label="positive"
    )
    neg_line = mlines.Line2D(
        [], [], color="gray", linestyle='--', label="negative"
    )
    pos_legend_line.append(pos_line)
    pos_legend_line.append(neg_line)
    rel_leg_ax.legend(
        handles=pos_legend_line, labelspacing=0.5, 
        title="Relation",frameon=False, framealpha=0.8, 
        title_fontproperties = {"weight":"bold", "size":"medium"}
    )
    
    degree_values = list(range(8, -1, -2))
    node_legend_handles = []
    for deg in degree_values:
        size = base_node_size + node_size_scale * deg
        node = plt.scatter(
            [], [], s=size, linewidths=0, 
            c=node_color, edgecolors="gray",
            alpha=0.9, label=f"{deg}"
        )
        node_legend_handles.append(node)
    deg_leg_ax.legend(
        handles=node_legend_handles, labelspacing=2, 
        title="Degree", frameon=False, framealpha=1,
        title_fontproperties = {"weight":"bold", "size":"medium"}
    )
    plt.tight_layout(w_pad=0.2)
    plt.show()
